import numpy as np
from scipy.linalg import solve_discrete_are, LinAlgError
import os

dir_current_file = os.path.dirname(__file__)  # 'Algorithm_demo_1\planner'
dir_parent_1 = os.path.dirname(dir_current_file)  # 'Algorithm_demo_1'

npz_path = os.path.join(dir_current_file, "all_results", "lqr_K", "lqr_K_table.npz")

# —— 请根据你的系统参数修改以下变量 ——
dt = 0.1        # 时间步长 Δt
L  = 5.6        # 轴距
# LQR 权重矩阵 Q, R，请替换为你的 self.lqr_Q 和 self.lqr_R
Q =  np.eye(5)
R = np.eye(2)

# 构建速度查表
v_min, v_max, v_step = -2.8, 12.0, 0.001
v_tbl = np.arange(v_min, v_max + v_step/2, v_step)  # 包含 v_max
K_tbl = np.zeros((len(v_tbl), 2, 5))



import numpy.linalg as la


def solve_dare(A, B, Q, R, max_iter=150, eps=1e-3):
    """
    迭代求解离散代数 Riccati 方程 (DARE)
    """

    x = Q
    x_next = Q
    for _ in range(max_iter):
        x_next = A.T @ x @ A - A.T @ x @ B @ \
                la.inv(R + B.T @ x @ B) @ B.T @ x @ A + Q
        if (abs(x_next - x)).max() < eps:
            break
        x = x_next
    return x


def dlqr(A, B, Q, R):
    """
    离散时间 LQR 求解
    x[k+1] = A x[k] + B u[k]
    cost = Σ (x.T Q x + u.T R u)
    Returns:
      K: 状态反馈增益 (2x5)
      X: Riccati 矩阵
      eig: 闭环特征值
    """
    X = solve_dare(A, B, Q, R)
    K = la.inv(B.T @ X @ B + R) @ (B.T @ X @ A)
    # eig_vals = la.eig(A - B @ K)[0]
    return K, X



# 离线计算每个速度下的最优增益，若失败则填 NaN
for i, v in enumerate(v_tbl):
    A = np.zeros((5, 5))
    A[0, 0] = 1.0
    A[0, 1] = dt
    A[1, 2] = v
    A[2, 2] = 1.0
    A[2, 3] = dt
    A[4, 4] = 1.0

    B = np.zeros((5, 2))
    B[3, 0] = v / L
    B[4, 1] = dt

    try:
        K , _ = dlqr(A,B,Q,R)
    except LinAlgError:
        K = np.full((2, 5), np.nan)

    K_tbl[i] = K

# 保存为 .npz 文件
np.savez(npz_path, v_tbl=v_tbl, K_tbl=K_tbl)
print(f"预计算并保存完毕: {npz_path}")
